#include <opencv2/opencv.hpp>
#include "posenet.h"
#include <windows.h>

cv::Mat getImage(const cv::Mat& im, cv::Size baseSize, float* scale) {
	int w = baseSize.width;
	int h = baseSize.height;
	int nh = h;
	float s = h / (float)im.rows;;
	int nw = im.cols * s;

	if (nw > w) {
		nw = w;
		s = w / (float)im.cols;
		nh = im.rows * s;
	}

	if (scale)*scale = 1 / s;
	cv::Rect dst(0, 0, nw, nh);
	cv::Mat bck = cv::Mat::zeros(h, w, CV_8UC3);
	resize(im, bck(dst), cv::Size(nw, nh));
	return bck;
}

void test0() {
	const char* prototxt = "../model/pose_deploy_linevec.prototxt";
	const char* caffemodel = "../model/pose_iter_440000.caffemodel";
	cv::Mat im = cv::imread("../misc/test_imgs/00.jpg");
	float toSize = 200;
	float width = float(im.cols);
	float height = float(im.rows);
	cv::Size baseSize = cv::Size(toSize, toSize / width * height); 
	printf("baseSize %d %d\n", baseSize.height, baseSize.width);
	cv::Mat raw_in;
	cv::resize(im, raw_in, baseSize);
	Posenet pnet(prototxt, caffemodel, baseSize);
	pnet.forward(raw_in);
}

void test1() {
	const char* prototxt = "../model/pose_deploy_linevec.prototxt";
	const char* caffemodel = "../model/pose_iter_440000.caffemodel";
	cv::Mat ori_img = cv::imread("../misc/test_imgs/00.jpg");
	float toSize = 200.;
	float width = float(ori_img.cols);
	float height = float(ori_img.rows);
	float scale = width / toSize;
	cv::Size baseSize = cv::Size(toSize, toSize / width * height);
	printf("baseSize %d %d\n", baseSize.height, baseSize.width);
	cv::Mat resz_img;
	cv::resize(ori_img, resz_img, baseSize);
	Posenet pnet(prototxt, caffemodel, baseSize);
	std::vector<float> keypoints;
	std::vector<int> keyshape;
	pnet.detect(resz_img, keypoints, keyshape);
	pnet.renderPoseKeypointsCpu(ori_img, keypoints, keyshape, 0.05, scale);
	cv::imshow("res", ori_img);
	cv::waitKey(0);
}

void test2() {
	const char* prototxt = "../model/pose_deploy_linevec.prototxt";
	const char* caffemodel = "../model/pose_iter_440000.caffemodel";
	cv::VideoCapture cap(0);
	float toSize = 200.;
	float width = 640.;
	float height = 480.;
	float scale = width / toSize;
	cv::Size baseSize = cv::Size(toSize, toSize / width * height);
	Posenet pnet(prototxt, caffemodel, baseSize);
	printf("baseSize %d %d\n", baseSize.height, baseSize.width);
	cv::Mat ori_img;
	std::vector<float> keypoints;
	std::vector<int> keyshape;
	while (cap.read(ori_img)){
		int64 t0 = cv::getTickCount();
		cv::Mat resz_img;
		cv::resize(ori_img, resz_img, baseSize);
		int64 t1 = cv::getTickCount();
		pnet.detect(resz_img, keypoints, keyshape);
		int64 t2 = cv::getTickCount();
		pnet.renderPoseKeypointsCpu(ori_img, keypoints, keyshape, 0.05, scale);
		int64 t3 = cv::getTickCount();
		printf("forward time used %fms\n", (t2 - t1) / cv::getTickFrequency() * 1000);
		printf("total time used %fms\n", (t3 - t0) / cv::getTickFrequency() * 1000);
		cv::imshow("res", ori_img);
		if (cv::waitKey(1) == VK_ESCAPE) {
			break;
		}
	}
}

int main() {
	TCHAR NPath[MAX_PATH];
	GetCurrentDirectory(MAX_PATH, NPath);
	std::cout << "current directory is " << NPath << "\n";
	if (caffe::GPUAvailable()) {
		printf("GPU is available\n");
		caffe::SetMode(caffe::GPU, 0);
	}
	else {
		printf("Only CPU is available\n");
	}
	test1();
	return 0;
}